#!/usr/bin/env bash

# Do "ps aux" on a set of nodes belonging to a single job, but exclude system processes.
# Author: Ole.H.Nielsen@fysik.dtu.dk
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/
# Usage: psjob jobid

# Requires ClusterShell for parallel commands,
# see https://wiki.fysik.dtu.dk/niflheim/SLURM
# and https://clustershell.readthedocs.io/en/latest/intro.html

# Possibly configure these variables:
SCONTROL="/usr/bin/scontrol"
CLUSH="/usr/bin/clush"
PS="/bin/ps"
PSFLAGS="-o pid,state,user,start,nlwp,cputime,%cpu,rssize,command --columns 100"

if test $# -ne 1
then
	echo ERROR. Usage: $0 jobid
	exit -1
fi
JOBID=$1

# System users excluded from the list
# The "68" is haldaemon on CentOS6 as printed by ps.
USERLIST="root rpc rpcuser daemon ntp smmsp sshd hpsmh named dbus 68 chrony polkitd munge"

# Make a deselect-list consisting only of existent users
deselect_list=""
sep=""
for u in $USERLIST
do
	if test -n "`getent passwd $u`"
	then
		# This user exists in the passwd database
		deselect_list="${deselect_list}${sep}${u}"
		sep=" "
	fi
done

# Temporary file
JOBSTATUS=/tmp/jobstatus.$$
# Catch signals
trap "rm -f $JOBSTATUS; exit 2" 1 2 3 14 15 19

# Check if this job-ID can be inquired successfully.
$SCONTROL show job $JOBID > $JOBSTATUS
if test "$?" != "0"
then
	echo Error inquiring about job $JOBID 
	exit 1
fi

# Check for job array
if test -n "`cat $JOBSTATUS | grep ArrayJobId`"
then
	ArrayJobId=`cat $JOBSTATUS | grep ArrayJobId | head -1 | awk '{print $2}' | sed /ArrayJobId=/s///`
	ArrayTaskId=`cat $JOBSTATUS | grep ArrayJobId | head -1 | awk '{print $3}' | sed /ArrayTaskId=/s///`
	if test "$JOBID" = "$ArrayJobId"
	then
		JOBID="${JOBID}_${ArrayTaskId}"
		echo Select JOBID=$JOBID
		# Call this command again with the modified JOBID
		$0 $JOBID
		exit 0
	fi
	echo "This job $JOBID is a job array: Head job $ArrayJobId task number $ArrayTaskId"
fi

# Check the job state
JOBSTATE="`cat $JOBSTATUS | grep JobState | awk '{print $1}'`"
if test "$JOBSTATE" != "JobState=RUNNING"
then
	echo The job $JOBID is not running, it has state=$JOBSTATE
	exit 1
fi

NODELIST="`cat $JOBSTATUS | grep '^   NodeList=' | awk -F= '{print $2}'`"
if test -z "$NODELIST"
then
	echo Error: The node list is empty
	exit 1
fi

# Get the number of nodes and node properties used
echo Nodelist for job-id $JOBID: $NODELIST
echo Node usage: `cat $JOBSTATUS | grep NumNodes= ` 
egrep 'RunTime|SubmitTime|StartTime' $JOBSTATUS

# Execute parallel shell on the job nodes
$CLUSH -b -w $NODELIST $PS $PSFLAGS --deselect -u \""$USERLIST"\"

rm -f $JOBSTATUS
